{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# WebSockets"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Websockets are a protocol for two-way, persistent communication between a client and server. This is different from HTTP, which uses a request/response model where the client sends a request and the server responds. With websockets, either party can send messages at any time, and the other party can respond.\n",
    "\n",
    "This allows for different applications to be built, including things like chat apps, live-updating dashboards, and real-time collaborative tools, which would require constant polling of the server for updates with HTTP.\n",
    "\n",
    "In FastHTML, you can create a websocket route using the `@app.ws` decorator. This decorator takes a route path, and optional `conn` and `disconn` parameters representing the `on_connect` and `on_disconnect` callbacks in websockets, respectively. The function decorated by `@app.ws` is the main function that is called when a message is received.\n",
    "\n",
    "Here's an example of a basic websocket route:\n",
    "```python\n",
    "@app.ws('/ws', conn=on_conn, disconn=on_disconn)\n",
    "async def on_message(msg:str, send):\n",
    "    await send(Div('Hello ' + msg, id='notifications'))\n",
    "    await send(Div('Goodbye ' + msg, id='notifications'))\n",
    "```\n",
    "\n",
    "The `on_message` function is the main function that is called when a message is received and can be named however you like. Similar to standard routes, the arguments to `on_message` are automatically parsed from the websocket payload for you, so you don't need to manually parse the message content. However, certain argument names are reserved for special purposes. Here are the most important ones:\n",
    "\n",
    "- `send` is a function that can be used to send text data to the client.\n",
    "- `data` is a dictionary containing the data sent by the client.\n",
    "- `ws` is a reference to the websocket object.\n",
    "\n",
    "For example, we can send a message to the client that just connected like this:\n",
    "```python\n",
    "async def on_conn(send):\n",
    "    await send(Div('Hello, world!'))\n",
    "```\n",
    "\n",
    "Or if we receive a message from the client, we can send a message back to them:\n",
    "```python\n",
    "@app.ws('/ws', conn=on_conn, disconn=on_disconn)\n",
    "async def on_message(msg:str, send):\n",
    "    await send(Div('You said: ' + msg, id='notifications'))\n",
    "    # or...\n",
    "    return Div('You said: ' + msg, id='notifications')\n",
    "```\n",
    "\n",
    "On the client side, we can use HTMX's websocket extension to open a websocket connection and send/receive messages. For example:\n",
    "```python\n",
    "from fasthtml.common import *\n",
    "\n",
    "app = FastHTML(exts='ws')\n",
    "\n",
    "@app.get('/')\n",
    "def home():\n",
    "    cts = Div(\n",
    "        Div(id='notifications'),\n",
    "        Form(Input(id='msg'), id='form', ws_send=True),\n",
    "        hx_ext='ws', ws_connect='/ws')\n",
    "    return Titled('Websocket Test', cts)\n",
    "```\n",
    "\n",
    "This will create a websocket connection to the server on route `/ws`, and send any form submissions to the server via the websocket. The server will then respond by sending a message back to the client. The client will then update the message div with the message from the server using Out of Band Swaps, which means that the content is swapped with the same id without reloading the page.\n",
    "\n",
    "::: {.callout-note}\n",
    "Make sure you set `exts='ws'` when creating your `FastHTML` object if you want to use websockets so the extension is loaded.\n",
    ":::\n",
    "\n",
    "Putting it all together, the code for the client and server should look like this:\n",
    "\n",
    "```python\n",
    "from fasthtml.common import *\n",
    "\n",
    "app = FastHTML(exts='ws')\n",
    "rt = app.route\n",
    "\n",
    "@rt('/')\n",
    "def get():\n",
    "    cts = Div(\n",
    "        Div(id='notifications'),\n",
    "        Form(Input(id='msg'), id='form', ws_send=True),\n",
    "        hx_ext='ws', ws_connect='/ws')\n",
    "    return Titled('Websocket Test', cts)\n",
    "\n",
    "@app.ws('/ws')\n",
    "async def ws(msg:str, send):\n",
    "    await send(Div('Hello ' + msg, id='notifications'))\n",
    "\n",
    "serve()\n",
    "```\n",
    "\n",
    "This is a fairly simple example and could be done just as easily with standard HTTP requests, but it illustrates the basic idea of how websockets work. Let's look at a more complex example next."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Session data in Websockets\n",
    "\n",
    "Session data is shared between standard HTTP routes and Websockets. This means you can access, for example, logged in user ID inside websocket handler:\n",
    "\n",
    "```python\n",
    "from fasthtml.common import *\n",
    "\n",
    "app = FastHTML(exts='ws')\n",
    "rt = app.route\n",
    "\n",
    "@rt('/login')\n",
    "def get(session):\n",
    "    session[\"person\"] = \"Bob\"\n",
    "    return \"ok\"\n",
    "\n",
    "@app.ws('/ws')\n",
    "async def ws(msg:str, send, session):\n",
    "    await send(Div(f'Hello {session.get(\"person\")}' + msg, id='notifications'))\n",
    "\n",
    "serve()\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Real-Time Chat App\n",
    "\n",
    "Let's put our new websocket knowledge to use by building a simple chat app. We will create a chat app where multiple users can send and receive messages in real time.\n",
    "\n",
    "Let's start by defining the app and the home page:\n",
    "```python\n",
    "from fasthtml.common import *\n",
    "\n",
    "app = FastHTML(exts='ws')\n",
    "rt = app.route\n",
    "\n",
    "msgs = []\n",
    "@rt('/')\n",
    "def home(): return Div(\n",
    "    Div(Ul(*[Li(m) for m in msgs], id='msg-list')),\n",
    "    Form(Input(id='msg'), id='form', ws_send=True),\n",
    "    hx_ext='ws', ws_connect='/ws')\n",
    "```\n",
    "\n",
    "Now, let's handle the websocket connection. We'll add a new route for this along with an `on_conn` and `on_disconn` function to keep track of the users currently connected to the websocket. Finally, we will handle the logic for sending messages to all connected users.\n",
    "\n",
    "```python\n",
    "users = {}\n",
    "def on_conn(ws, send): users[str(id(ws))] = send\n",
    "def on_disconn(ws): users.pop(str(id(ws)), None)\n",
    "\n",
    "@app.ws('/ws', conn=on_conn, disconn=on_disconn)\n",
    "async def ws(msg:str):\n",
    "    msgs.append(msg)\n",
    "    # Use associated `send` function to send message to each user\n",
    "    for u in users.values(): await u(Ul(*[Li(m) for m in msgs], id='msg-list'))\n",
    "\n",
    "serve()\n",
    "```\n",
    "\n",
    "We can now run this app with `python chat_ws.py` and open multiple browser tabs to `http://localhost:5001`. You should be able to send messages in one tab and see them appear in the other tabs."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### A Work in Progress"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This page (and Websocket support in FastHTML) is a work in progress. Questions, PRs, and feedback are welcome!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
